#include "nnslitemtree.h"

static NNSLITE_RET nnsliteMTreeNodeInit(nnsliteMTreeNode* node, nnsliteMTree* tree);
static NNSLITE_RET nnsliteMTreeNodeFree(nnsliteMTreeNode* node);
static NNSLITE_RET nnsliteMTreeNodeWrite(nnsliteMTreeNode* node,FILE* fp);
static NNSLITE_RET nnsliteMTreeNodeRead(nnsliteMTreeNode* node, FILE* fp);
static NNSLITE_RET nnsliteMTreeNodeFindSplitSeeds(nnsliteMTree* tree, u32 nodeindex, nnsliteObject* seed1, nnsliteObject* seed2);
static NNSLITE_RET nnsliteMTreeNodeAbsorb(nnsliteMTree* tree, u32 nodeindex, nnsliteMTreeNode* seedNode1, nnsliteMTreeNode* seedNode2);
static NNSLITE_RET nnsliteMTreeNodeSplit(nnsliteMTree* tree, u32 nodeindex);
static NNSLITE_RET nnsliteMTreeNodeInsert(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object);
static NNSLITE_RET nnsliteMTreeNodeNNSearch(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object, u32 top, nnsliteObject* results, f64* resultDists);
static NNSLITE_RET nnsliteMTreeNodePrint(nnsliteMTreeNode* node, FILE* fp, u32 indent);
static NNSLITE_RET nnsliteMTreeRemoveNode(nnsliteMTree* tree, u32 nodeindex);

/** nnsliteMTree  */
static NNSLITE_RET nnsliteMTreeNodeInit(nnsliteMTreeNode* node, nnsliteMTree* tree)
{
	node->isLeaf = TRUE;
	node->pTree = tree;
	node->nodeIndex = tree->nodes.size;
	node->parentIndex = U32_NULL;
	node->childCount = 0;
	node->depth = 0;
	node->radius = 0.0;
	node->numObjectInPage = 0 ;
	nnsliteObjectInit(&node->object);
	nnsliteDataPageInit(&node->dataPage);
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeFree(nnsliteMTreeNode* node)
{
	nnsliteObjectFree(&node->object);
	node->numObjectInPage = 0;
	nnsliteDataPageFree(&node->dataPage);
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeWrite(nnsliteMTreeNode* node,FILE* fp)
{
	fwrite(&node->nodeIndex, sizeof(u32),1,fp);
	fwrite(node->childIndexes, sizeof(u32)*NNSLITE_MTREE_MAX_CHILD, 1,fp);
	fwrite(&node->childCount, sizeof(u32), 1, fp);
	fwrite(&node->isLeaf, sizeof(u8), 1, fp);
	fwrite(&node->depth, sizeof(u32), 1, fp);
	fwrite(&node->radius, sizeof(f64), 1, fp);
	fwrite(&node->object.id, sizeof(u64), 1, fp);
	fwrite(&node->object.size, sizeof(u32), 1, fp);
	fwrite(node->object.data, node->object.size, 1, fp);
	if (node->isLeaf)
	{
		nnsliteDataPageWriteTo(&node->dataPage, fp);
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeRead(nnsliteMTreeNode* node, FILE* fp)
{
	fread(&node->nodeIndex, sizeof(u32),1,fp);
	fread(node->childIndexes, sizeof(u32)*NNSLITE_MTREE_MAX_CHILD, 1, fp);
	fread(&node->childCount, sizeof(u32), 1, fp);
	fread(&node->isLeaf, sizeof(u8), 1, fp);
	fread(&node->depth, sizeof(u32), 1, fp);
	fread(&node->radius, sizeof(f64), 1, fp);
	fread(&node->object.id, sizeof(u64), 1, fp);
	fread(&node->object.size, sizeof(u32), 1, fp);
	if (node->object.data != NULL)
		free(node->object.data);
	node->object.data = (u8*)malloc(node->object.size);
	fread(node->object.data, node->object.size, 1, fp);
	if (node->isLeaf)
	{
		nnsliteDataPageReadFrom(&node->dataPage, fp);
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeFindSplitSeeds(nnsliteMTree* tree, u32 nodeindex, nnsliteObject* seed1, nnsliteObject* seed2)
{
	s32 ret;
	f64 maxDist = 0.0;
	f64 dist;
	nnsliteMTreeNode node;
	nnsliteVectorGet(&tree->nodes, nodeindex, &node);
	if (node.isLeaf)
	{
		u32 i;
		nnsliteObject obj;
		nnsliteDataPageCursor cur;

		ret = nnsliteDataPageCreateCursor(&node.dataPage, &cur);
		NNSLITE_ON_FAILED_RETURN(ret);

		nnsliteObjectInit(&obj);
		for (i=0;i<node.numObjectInPage; i++)
		{
			ret = nnsliteDataPageNextObject(&cur, &obj);
			NNSLITE_ON_FAILED_RETURN(ret);
			dist = tree->distfunc(&node.object, &obj);
			if (dist > maxDist)
			{
				*seed2 = nnsliteObjectCopy(&obj);
			}
		}
		*seed1 = nnsliteObjectCopy(&node.object);

		nnsliteObjectFree(&obj);
	}
	else
	{
		nnsliteMTreeNode child;
		u32 childindex;
		u32 i;
		for(i=0;i<node.childCount;i++)
		{
			childindex = node.childIndexes[i];
			nnsliteVectorGet(&tree->nodes, childindex, &child);
			dist = tree->distfunc(&node.object, &child.object);
			if (dist> maxDist)
			{
				*seed2 = nnsliteObjectCopy(&child.object);
			}
		}
		*seed1 = nnsliteObjectCopy(&node.object);
	}

	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeAbsorb(nnsliteMTree* tree, u32 nodeindex, nnsliteMTreeNode* seedNode1, nnsliteMTreeNode* seedNode2)
{
	s32 ret;
	nnsliteMTreeNode parentNode;
	nnsliteObject obj;
	nnsliteObject* seed1 = &seedNode1->object;
	nnsliteObject* seed2 = &seedNode2->object;
	f64 dist1,dist2;
	
	nnsliteVectorGet(&tree->nodes, nodeindex, &parentNode);
	if (parentNode.isLeaf)
	{
		nnsliteDataPageCursor cur;

		ret = nnsliteObjectInit(&obj);
		NNSLITE_ON_FAILED_RETURN(ret);

		ret = nnsliteDataPageCreateCursor(&parentNode.dataPage, &cur);
		NNSLITE_ON_FAILED_RETURN(ret);

		while(nnsliteDataPageCursorIsEnd(&cur) == NNSLITE_RET_FALSE)
		{
			nnsliteDataPageNextObject(&cur, &obj);
			dist1 = tree->distfunc(&obj, seed1);
			dist2 = tree->distfunc(&obj, seed2);
			if (dist1 < dist2)
			{
				nnsliteDataPageWriteObj(&seedNode1->dataPage, &obj);
				seedNode1->numObjectInPage++;
				seedNode1->radius = MAX(seedNode1->radius, dist1);
			}
			else
			{
				nnsliteDataPageWriteObj(&seedNode2->dataPage, &obj);
				seedNode2->numObjectInPage++;
				seedNode2->radius = MAX(seedNode2->radius, dist2);
			}
		}
		nnsliteObjectFree(&obj);
	}
	else
	{
		u32 i;
		seedNode1->radius = 0.0;
		seedNode2->radius = 0.0;
		for (i=0;i<parentNode.childCount; i++)
		{
			nnsliteMTreeNode childNode;
			u32 childIndex = parentNode.childIndexes[i];
			nnsliteVectorGet(&tree->nodes,childIndex, &childNode);
			dist1 = tree->distfunc(&childNode.object, seed1);
			dist2 = tree->distfunc(&childNode.object, seed2);
			if (dist1 < dist2)
			{
				seedNode1->childIndexes[seedNode1->childCount++] = childIndex;
				seedNode1->radius = MAX(seedNode1->radius, dist1+childNode.radius);
				childNode.parentIndex = seedNode1->nodeIndex;
			}
			else
			{
				seedNode2->childIndexes[seedNode2->childCount++] = childIndex;
				seedNode2->radius = MAX(seedNode2->radius, dist2+childNode.radius);
				childNode.parentIndex = seedNode2->nodeIndex;
			}
			nnsliteVectorSet(&tree->nodes,childIndex,&childNode);
		}
	}

	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeSplit(nnsliteMTree* tree, u32 nodeindex)
{
	s32 ret;
	nnsliteObject seed1, seed2;
	nnsliteMTreeNode seedNode1;
	nnsliteMTreeNode seedNode2;
	nnsliteMTreeNode node;
	nnsliteMTreeNode parentNode;


	/* get the seeds to split */
	nnsliteObjectInit(&seed1);
	nnsliteObjectInit(&seed2);
	ret = nnsliteMTreeNodeFindSplitSeeds(tree, nodeindex, &seed1, &seed2);
	NNSLITE_ON_FAILED_RETURN(ret);

	/* get the node*/
	nnsliteVectorGet(&tree->nodes, nodeindex, &node);

	/** create two new nodes */
	nnsliteMTreeNodeInit(&seedNode1, tree);
	nnsliteMTreeNodeInit(&seedNode2, tree);
	seedNode1.object = nnsliteObjectCopy(&seed1);
	seedNode2.object = nnsliteObjectCopy(&seed2);
	seedNode1.depth = node.depth;
	seedNode2.depth = node.depth;
	seedNode1.isLeaf = node.isLeaf;
	seedNode2.isLeaf = node.isLeaf;
	seedNode1.nodeIndex = node.nodeIndex;
	seedNode2.nodeIndex = tree->nodes.size;

	/* absorbs the nodes by the seeds */
	ret = nnsliteMTreeNodeAbsorb(tree, nodeindex, &seedNode1, &seedNode2);
	NNSLITE_ON_FAILED_RETURN(ret);

	if (node.parentIndex != U32_NULL) /* has parent node */
	{
		u32 originalNodeIndex = node.nodeIndex;
		nnsliteVectorGet(&tree->nodes, node.parentIndex, &parentNode);
		parentNode.childIndexes[parentNode.childCount++] = seedNode2.nodeIndex;
		seedNode1.parentIndex = node.parentIndex;
		seedNode2.parentIndex = node.parentIndex;
		/* add the two seed nodes */
		nnsliteMTreeNodeFree(&node);
		nnsliteVectorSet(&tree->nodes, originalNodeIndex, &seedNode1);
		nnsliteVectorPushback(&tree->nodes, &seedNode2);
		nnsliteVectorSet(&tree->nodes,parentNode.nodeIndex, &parentNode);
		/* check the condition for the splitting of parent node */
		if (parentNode.childCount == NNSLITE_MTREE_MAX_CHILD)
		{
			ret = nnsliteMTreeNodeSplit(tree, parentNode.nodeIndex);
			NNSLITE_ON_FAILED_RETURN(ret);
		}
	}
	else /* does not have parent node */
	{
		u32 i;
		u32 originalNodeIndex = node.nodeIndex;
		nnsliteMTreeNode newRootNode;
		/* update the depth of all nodes */
		for(i=0;i<tree->nodes.size;i++)
		{
			nnsliteMTreeNode eachNode;
			nnsliteVectorGet(&tree->nodes, i, &eachNode);
			eachNode.depth++;
			nnsliteVectorSet(&tree->nodes, i, &eachNode);
		}
		seedNode1.depth++;
		seedNode2.depth++;
		/* initialize the new root node */
		nnsliteMTreeNodeInit(&newRootNode, tree);
		newRootNode.depth = 0;
		newRootNode.isLeaf = FALSE;
		newRootNode.nodeIndex = tree->nodes.size+1;
		newRootNode.childIndexes[newRootNode.childCount++] = seedNode1.nodeIndex;
		newRootNode.childIndexes[newRootNode.childCount++] = seedNode2.nodeIndex;
		newRootNode.object = nnsliteObjectCopy(&node.object);
		newRootNode.radius = node.radius;
		/* add the two seed nodes */
		seedNode1.parentIndex = newRootNode.nodeIndex;
		seedNode2.parentIndex = newRootNode.nodeIndex;
		nnsliteMTreeNodeFree(&node);
		nnsliteVectorSet(&tree->nodes, originalNodeIndex, &seedNode1);
		nnsliteVectorPushback(&tree->nodes, &seedNode2);		
		/* add the new root node */
		nnsliteVectorPushback(&tree->nodes, &newRootNode);
		/* update the new root node index of the tree */
		tree->rootNodeIndex = newRootNode.nodeIndex;
	}
	
	nnsliteObjectFree(&seed1);
	nnsliteObjectFree(&seed2);
	
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeInsert(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object)
{
	s32 ret;
	f64 distToRoute;
	nnsliteMTreeNode node;
	nnsliteDataPage* page;
	nnsliteDistFunc distfunc = tree->distfunc;

	/* get the node*/
	nnsliteVectorGet(&tree->nodes, nodeindex, &node);	
	if (NNSLITE_OBJ_ISNULL(&node.object))
	{
		node.object = nnsliteObjectCopy(object);
		node.radius = 0.0;
	}
	
	distToRoute = distfunc(&node.object, object);
	node.radius = MAX(node.radius, distToRoute);
	
	if (node.isLeaf)
	{
		page = &node.dataPage;
		nnsliteDataPageWriteObj(page, object);
		node.numObjectInPage++;

		nnsliteVectorSet(&tree->nodes, nodeindex, &node);

		if (node.numObjectInPage >= NNSLITE_MTREE_MAX_CHILD && node.depth < tree->maxDepth)
		{
			ret = nnsliteMTreeNodeSplit(tree, nodeindex);
			NNSLITE_ON_FAILED_RETURN(ret);
		}
	}
	else
	{
		f64 d;
		f64 minDist = (f64)DBL_MAX;
		u32 i;
		nnsliteMTreeNode childNode;
		nnsliteMTreeNode* nearestChildNode = NULL;
		for (i=0; i<node.childCount; i++)
		{
			u32 childIndex = node.childIndexes[i];
			nnsliteVectorGet(&tree->nodes, childIndex, &childNode);
			d = distfunc(&childNode.object, object);
			if (d < minDist)
			{
				minDist = d;
				nearestChildNode = &childNode;
			}
		}

		nnsliteVectorSet(&tree->nodes, nodeindex, &node);
		ret = nnsliteMTreeNodeInsert(tree, nearestChildNode->nodeIndex, object);
		NNSLITE_ON_FAILED_RETURN(ret);
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodeNNSearch(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object, u32 top, nnsliteObject* results,f64* resultDists)
{
	s32 ret;
	u32 i;
	f64 dist;
	nnsliteMTreeNode node;

	nnsliteVectorGet(&tree->nodes, nodeindex, &node);
	if (node.isLeaf)
	{
		nnsliteObject obj;
		nnsliteDataPageCursor cur;
		s32 j;
		u32 pos;

		ret = nnsliteObjectInit(&obj);
		NNSLITE_ON_FAILED_RETURN(ret);

		ret = nnsliteDataPageCreateCursor(&node.dataPage, &cur);
		NNSLITE_ON_FAILED_RETURN(ret);

		for(i=0;i<node.numObjectInPage; i++)
		{
			nnsliteDataPageNextObject(&cur, &obj);
			dist = tree->distfunc(object, &obj);
			for (j=top-1;j>=0;j--)
			{
				if (dist > resultDists[j])
					break;
			}
			pos = (u32)(j+1);
			if (pos < top)
			{
				/* An example : 0 1 2 3 4  top=5,j=2 */
				nnsliteObjectFree(&results[top-1]); /* Tick out the last one */
				if (pos < top-1)
				{
					memmove(&results[pos+1],&results[pos], (top-pos-1)*sizeof(nnsliteObject));
					memmove(&resultDists[pos+1], &resultDists[pos], (top-pos-1)*sizeof(f64));
				}
				results[pos] = nnsliteObjectCopy(&obj);
				resultDists[pos] = dist;
			}
		}

		nnsliteObjectFree(&obj);
	}
	else
	{
		nnsliteMTreeNode childNode;
		u32 childIndex;
		for (i=0; i<node.childCount; i++)
		{
			childIndex = node.childIndexes[i];
			nnsliteVectorGet(&tree->nodes, childIndex, &childNode);
			dist = tree->distfunc(object, &childNode.object);
			if (dist - childNode.radius < resultDists[top-1])
			{
				ret = nnsliteMTreeNodeNNSearch(tree, childIndex, object, top, results, resultDists);
				NNSLITE_ON_FAILED_RETURN(ret);
			}
		}
	}

	return NNSLITE_RET_OK;
}


static NNSLITE_RET nnsliteMTreeRemoveNode(nnsliteMTree* tree, u32 indexnode)
{
	u32 i,j;
	nnsliteMTreeNode node;
	for (i=0;i<tree->nodes.size; i++)
	{
		nnsliteVectorGet(&tree->nodes, i, &node);
		for(j=0;j<node.childCount;j++)
		{
			if(node.childIndexes[j]>=indexnode)
				node.childIndexes[j]--;
		}
		if(node.nodeIndex>indexnode)
			node.nodeIndex--;
		if(node.parentIndex>indexnode)
			node.parentIndex--;
		nnsliteVectorSet(&tree->nodes, i, &node);
	}
	if (tree->rootNodeIndex >= indexnode)
		tree->rootNodeIndex--;
	return nnsliteVectorDelete(&tree->nodes, indexnode);
}

NNSLITE_RET nnsliteMTreeInit(nnsliteMTree* tree, nnsliteDistFunc func)
{
	s32 ret;
	nnsliteMTreeNode root;
	if (NNSLITE_RET_OK != nnsliteVectorInit(&tree->nodes, sizeof(nnsliteMTreeNode)))
		return NNSLITE_RET_ERR;

	ret = nnsliteMTreeNodeInit(&root, tree);
	NNSLITE_ON_FAILED_RETURN(ret);
	root.depth = 0;
	root.radius = 0.0;
	nnsliteVectorPushback(&tree->nodes, (void*)(&root));

	tree->distfunc = func;
	tree->maxDepth = NNSLITE_MTREE_MAX_DEPTH;
	tree->rootNodeIndex = 0;
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeFree(nnsliteMTree* tree)
{
	u32 i;
	for (i=0;i<tree->nodes.size;i++)
	{
		nnsliteMTreeNode node;
		nnsliteVectorGet(&tree->nodes, i, &node);
		nnsliteMTreeNodeFree(&node);
	}
	nnsliteVectorFree(&tree->nodes);
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeInsert(nnsliteMTree* tree, const nnsliteObject* object)
{
	return nnsliteMTreeNodeInsert(tree, tree->rootNodeIndex, object);
}

NNSLITE_RET nnsliteMTreeNNSearch(nnsliteMTree* tree, const nnsliteObject* object, u32 top,nnsliteObject* results)
{
	f64* resultDists = (f64*)malloc(sizeof(f64)*top);
	u32 i;
	s32 ret;
	
	for (i=0; i<top; i++)
		resultDists[i] = (f64)F64_MAX_VAL;
	
	ret = nnsliteMTreeNodeNNSearch(tree, tree->rootNodeIndex, object, top, results, resultDists);
	NNSLITE_ON_FAILED_RETURN(ret);

	free(resultDists);
	return NNSLITE_RET_OK;
}


NNSLITE_RET nnsliteMTreeWrite(nnsliteMTree* tree, FILE* fp)
{
	u32 i;
	fwrite(&tree->nodes.size, sizeof(u32),1, fp);
	for (i=0;i<tree->nodes.size;i++)
	{
		nnsliteMTreeNode node;
		nnsliteVectorGet(&tree->nodes, i, &node);
		nnsliteMTreeNodeWrite(&node, fp);
	}
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeRead(nnsliteMTree* tree, FILE* fp)
{
	u32 size,i;

	nnsliteMTreeNode node;
	nnsliteVectorClear(&tree->nodes);
	
	fread(&size, sizeof(u32), 1, fp);
	for(i=0;i<size;i++)
	{
		nnsliteMTreeNodeInit(&node, tree);
		nnsliteMTreeNodeRead(&node, fp);
		nnsliteVectorPushback(&tree->nodes, (void*)(&node));
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET nnsliteMTreeNodePrint(nnsliteMTreeNode* node, FILE* fp, u32 indent)
{
	nnsliteMTree* tree = node->pTree;
	nnsliteMTreeNode childNode;
	u32 childindex;
	u32 i;

	for (i=0;i<indent;i++)
		fprintf(fp,"\t");
	fprintf(fp,"Node ObjectID=%ld, ",node->object.id);
	fprintf(fp,"Radius=%lf, ",node->radius);
	fprintf(fp,"ChildCount=%d, ", node->childCount);

	if (node->isLeaf)
	{
		nnsliteObject obj;
		nnsliteDataPageCursor cur;
		
		nnsliteObjectInit(&obj);
		nnsliteDataPageCreateCursor(&node->dataPage, &cur);
		fprintf(fp, "DataPage (Count=%d): ",node->numObjectInPage);
		for (i=0;i<node->numObjectInPage;i++)
		{
			nnsliteDataPageNextObject(&cur, &obj);
			fprintf(fp,"%ld",obj.id);
			if (i<node->numObjectInPage-1)
				fprintf(fp,", ");
		}
		fprintf(fp,"\n");
		nnsliteObjectFree(&obj);
	}
	else
	{
		fprintf(fp,"\n");
		for (i=0; i<node->childCount; i++)
		{
			childindex = node->childIndexes[i];
			nnsliteVectorGet(&tree->nodes, childindex, &childNode);
			nnsliteMTreeNodePrint(&childNode, fp, indent+1);
		}
	}
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreePrint(nnsliteMTree* tree, FILE* fp)
{
	nnsliteMTreeNode rootNode;
	nnsliteVectorGet(&tree->nodes, tree->rootNodeIndex, &rootNode);
	return nnsliteMTreeNodePrint(&rootNode, fp, 0);
}
